import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from utils.csv_hanlder import select_and_read_files
from utils.twiss_hanlder import calculate_twiss_parameters
from utils.console_hanlder import console


def plot_beam_size_vs_frequency_log(dfs, filenames=None, save_path='result/beam_size_vs_frequency_logscale.png'):
    """
    Plot RMS beam size in x and y directions vs frequency
    """
    if not dfs:
        print("没有数据可绘制。")
        return
    
    frequencies = []
    x_beam_sizes = []  # x-direction beam size (RMS)
    y_beam_sizes = []  # y-direction beam size (RMS)
    x_emittances = []  # x-direction emittance
    y_emittances = []  # y-direction emittance
    x_beta_values = [] # x-direction beta function
    y_beta_values = [] # y-direction beta function
    x_alpha_values = [] # x-direction alpha parameter
    y_alpha_values = [] # y-direction alpha parameter
    
    # Collect all data first
    data_list = []
    for idx, df in enumerate(dfs):
        f_val = df['f'].iloc[0]
        x = df['x']
        y = df['y']
        vx = df['vx']
        vy = df['vy']
        vz = df['vz']
        
        # Calculate x', y'
        vz_nonzero = np.where(vz != 0, vz, 1e-10)
        x_prime = (vx / vz_nonzero)
        y_prime = (vy / vz_nonzero)
        
        # Calculate Twiss parameters
        beta_x, alpha_x, gamma_x, emittance_x = calculate_twiss_parameters(x, x_prime)
        beta_y, alpha_y, gamma_y, emittance_y = calculate_twiss_parameters(y, y_prime)
        
        # Calculate RMS beam size: σ_x = √(εβ)
        beam_size_x = np.sqrt(emittance_x * beta_x)
        beam_size_y = np.sqrt(emittance_y * beta_y)
        
        data_list.append({
            'freq': f_val,
            'beam_size_x': beam_size_x,
            'beam_size_y': beam_size_y,
            'emittance_x': emittance_x,
            'emittance_y': emittance_y,
            'beta_x': beta_x,
            'beta_y': beta_y,
            'alpha_x': alpha_x,
            'alpha_y': alpha_y
        })
    
    # Sort by frequency to match emmitance.py behavior
    data_list.sort(key=lambda x: x['freq'])
    
    # Extract sorted data
    for data in data_list:
        frequencies.append(data['freq'])
        x_beam_sizes.append(data['beam_size_x'])
        y_beam_sizes.append(data['beam_size_y'])
        x_emittances.append(data['emittance_x'])
        y_emittances.append(data['emittance_y'])
        x_beta_values.append(data['beta_x'])
        y_beta_values.append(data['beta_y'])
        x_alpha_values.append(data['alpha_x'])
        y_alpha_values.append(data['alpha_y'])
    
    # Create figure
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()
    
    # Convert to numpy arrays for easier handling
    frequencies = np.array(frequencies)
    x_beam_sizes = np.array(x_beam_sizes)
    y_beam_sizes = np.array(y_beam_sizes)
    x_emittances = np.array(x_emittances)
    y_emittances = np.array(y_emittances)
    x_beta_values = np.array(x_beta_values)
    y_beta_values = np.array(y_beta_values)
    x_alpha_values = np.array(x_alpha_values)
    y_alpha_values = np.array(y_alpha_values)
    
    # 1. Beam size vs frequency
    scatter_x = axes[0].scatter(frequencies, x_beam_sizes, c=frequencies, cmap='plasma', s=50, alpha=0.5, label='X-direction')
    scatter_y = axes[0].scatter(frequencies, y_beam_sizes, c=frequencies, cmap='viridis', s=50, alpha=0.5, label='Y-direction')
    plt.colorbar(scatter_x, ax=axes[0], shrink=0.8, label='Frequency [Hz]')
    axes[0].set_xlabel('Frequency [Hz]')
    axes[0].set_ylabel('RMS Beam Size [mm] (log scale)')
    axes[0].set_yscale('log')  # 设置对数尺度
    axes[0].set_title('Beam Size vs Frequency')
    axes[0].grid(True, alpha=0.3)
    axes[0].legend()
    
    # 2. Emittance vs frequency
    scatter_x = axes[1].scatter(frequencies, x_emittances, c=frequencies, cmap='plasma', s=50, alpha=0.5, label='X-direction')
    scatter_y = axes[1].scatter(frequencies, y_emittances, c=frequencies, cmap='viridis', s=50, alpha=0.5, label='Y-direction')
    plt.colorbar(scatter_x, ax=axes[1], shrink=0.8, label='Frequency [Hz]')
    axes[1].set_xlabel('Frequency [Hz]')
    axes[1].set_ylabel('Emittance [mm·rad] (log scale)')
    axes[1].set_yscale('log')  # 设置对数尺度
    axes[1].set_title('Emittance vs Frequency')
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    # 3. Beta function vs frequency
    axes[2].scatter(frequencies, x_beta_values, c='blue', s=50, alpha=0.5, label='β_x')
    axes[2].scatter(frequencies, y_beta_values, c='red', s=50, alpha=0.5, label='β_y')
    axes[2].set_xlabel('Frequency [Hz]')
    axes[2].set_ylabel('Beta Function [mm/rad] (log scale)')
    axes[2].set_yscale('log')  # 设置对数尺度
    axes[2].set_title('Beta Function vs Frequency')
    axes[2].grid(True, alpha=0.3)
    axes[2].legend()
    
    # 4. Alpha parameter vs frequency
    axes[3].scatter(frequencies, x_alpha_values, c='blue', s=50, alpha=0.5, label='α_x')
    axes[3].scatter(frequencies, y_alpha_values, c='red', s=50, alpha=0.5, label='α_y')
    axes[3].set_xlabel('Frequency [Hz]')
    axes[3].set_ylabel('Alpha Parameter')
    axes[3].set_title('Alpha Parameter vs Frequency')
    axes[3].grid(True, alpha=0.3)
    axes[3].axhline(y=0, color='k', linestyle='--', alpha=0.5)
    axes[3].legend()
    
    # 5. Beam size ratio (asymmetry)
    size_ratio = np.array(x_beam_sizes) / np.array(y_beam_sizes)
    scatter_ratio = axes[4].scatter(frequencies, size_ratio, c=frequencies, cmap='plasma', s=50, alpha=0.5)
    plt.colorbar(scatter_ratio, ax=axes[4], shrink=0.8, label='Frequency [Hz]')
    axes[4].set_xlabel('Frequency [Hz]')
    axes[4].set_ylabel('X/Y Beam Size Ratio')
    axes[4].set_title('Beam Asymmetry (X/Y Ratio)')
    axes[4].grid(True, alpha=0.3)
    axes[4].axhline(y=1, color='r', linestyle='--', alpha=0.5, label='Symmetry line')
    axes[4].legend()
    
    # 6. Product of beam sizes (approximate envelope area)
    envelope_area = np.array(x_beam_sizes) * np.array(y_beam_sizes)
    scatter_env = axes[5].scatter(frequencies, envelope_area, c=frequencies, cmap='plasma', s=50, alpha=0.5)
    plt.colorbar(scatter_env, ax=axes[5], shrink=0.8, label='Frequency [Hz]')
    axes[5].set_xlabel('Frequency [Hz]')
    axes[5].set_ylabel('X×Y Product [mm²] (log scale)')
    axes[5].set_yscale('log')  # 设置对数尺度
    axes[5].set_title('Beam Envelope Area (X×Y)')
    axes[5].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Beam size analysis plot saved to: {save_path}")
    plt.close(fig)
    
    # Log comprehensive statistics
    console.log_section("束流尺寸和 TWISS 参数分析")
    
    summary_rows = []
    for direction, sizes, emittances, betas, alphas, freqs in [
        ('X', x_beam_sizes, x_emittances, x_beta_values, x_alpha_values, frequencies),
        ('Y', y_beam_sizes, y_emittances, y_beta_values, y_alpha_values, frequencies)
    ]:
        min_idx = np.argmin(sizes)
        min_size = sizes[min_idx]
        freq_at_min = freqs[min_idx]
        emit_at_min = emittances[min_idx]
        beta_at_min = betas[min_idx]
        alpha_at_min = alphas[min_idx]
        
        # Determine convergence status
        if alpha_at_min > 0.01:
            status = "收敛"
        elif alpha_at_min < -0.01:
            status = "发散"
        else:
            status = "束腰"
        
        summary_rows.append([
            direction, f"{min_size:.2e}", str(freq_at_min), f"{emit_at_min:.2e}", 
            f"{beta_at_min:.2e}", f"{alpha_at_min:.2e}", status
        ])
    
    console.log_table(
        ['方向', '最小尺寸', '最小频率', '发射度', 'Beta', 'Alpha', '状态'],
        summary_rows
    )
    
    # Find frequency with minimum envelope area
    envelope_area = np.array(x_beam_sizes) * np.array(y_beam_sizes)
    min_area_idx = np.argmin(envelope_area)
    console.log(f"最小束流包络面积: {envelope_area[min_area_idx]:.2e} mm² 在频率 f={frequencies[min_area_idx]} Hz")
    
    # Find alpha local minima (points where alpha is close to zero)
    console.log_detailed_results("\n" + "=" * 60)
    console.log_detailed_results("Alpha 接近 0 的数据点 (局部最小值)")
    console.log_detailed_results("=" * 60)
    
    def find_alpha_minima(freqs, alphas, direction_name):
        """Find local minima of |alpha| (points where alpha is close to zero)"""
        freqs = np.array(freqs)
        alphas = np.array(alphas)
        abs_alphas = np.abs(alphas)
        
        # Find local minima of |alpha| using find_peaks on -abs_alphas
        # prominence should be relative to the range of abs_alphas
        prominence = 0.01 * np.ptp(abs_alphas)  # 1% of the range
        distance = len(alphas) // 100  # Minimum distance between peaks (1% of data points)
        
        # Find peaks in -abs_alphas (which are minima in abs_alphas)
        minima_indices, properties = find_peaks(-abs_alphas, prominence=prominence, distance=distance)
        
        # Also include points where |alpha| < threshold
        alpha_threshold = 0.001
        near_zero_indices = np.where(abs_alphas < alpha_threshold)[0]
        
        # Combine and remove duplicates
        all_indices = set(minima_indices) | set(near_zero_indices)
        
        # Create list of (freq, alpha) pairs, but filter by threshold
        all_points = []
        for idx in all_indices:
            if abs_alphas[idx] < alpha_threshold:
                all_points.append((freqs[idx], alphas[idx]))
        
        # Sort by frequency
        all_points.sort(key=lambda x: x[0])
        
        return all_points
    
    # X direction
    x_zero_points = find_alpha_minima(frequencies, x_alpha_values, "X")
    if x_zero_points:
        console.log_detailed_results(f"\nX 方向 (共 {len(x_zero_points)} 个数据点):")
        console.log_detailed_results(f"{'频率 (Hz)':<20} {'Alpha':<15}")
        console.log_detailed_results("-" * 60)
        for freq, alpha in x_zero_points:
            console.log_detailed_results(f"{freq:<20.0f} {alpha:<15.6e}")
    else:
        console.log_detailed_results("\nX 方向: 未找到 Alpha 接近 0 的数据点")
    
    # Y direction
    y_zero_points = find_alpha_minima(frequencies, y_alpha_values, "Y")
    if y_zero_points:
        console.log_detailed_results(f"\nY 方向 (共 {len(y_zero_points)} 个数据点):")
        console.log_detailed_results(f"{'频率 (Hz)':<20} {'Alpha':<15}")
        console.log_detailed_results("-" * 60)
        for freq, alpha in y_zero_points:
            console.log_detailed_results(f"{freq:<20.0f} {alpha:<15.6e}")
    else:
        console.log_detailed_results("\nY 方向: 未找到 Alpha 接近 0 的数据点")
    
    console.log_detailed_results("=" * 60)

def plot_beam_size_vs_frequency_linear(dfs, filenames=None, save_path='result/beam_size_vs_frequency_linear.png'):
    """
    Plot RMS beam size in x and y directions vs frequency using LINEAR scale
    """
    if not dfs:
        print("没有数据可绘制。")
        return
    
    frequencies = []
    x_beam_sizes = []  # x-direction beam size (RMS)
    y_beam_sizes = []  # y-direction beam size (RMS)
    x_emittances = []  # x-direction emittance
    y_emittances = []  # y-direction emittance
    x_beta_values = [] # x-direction beta function
    y_beta_values = [] # y-direction beta function
    x_alpha_values = [] # x-direction alpha parameter
    y_alpha_values = [] # y-direction alpha parameter
    
    # Collect all data first
    data_list = []
    for idx, df in enumerate(dfs):
        f_val = df['f'].iloc[0]
        x = df['x']
        y = df['y']
        vx = df['vx']
        vy = df['vy']
        vz = df['vz']
        
        # Calculate x', y'
        vz_nonzero = np.where(vz != 0, vz, 1e-10)
        x_prime = (vx / vz_nonzero)
        y_prime = (vy / vz_nonzero)
        
        # Calculate Twiss parameters
        beta_x, alpha_x, gamma_x, emittance_x = calculate_twiss_parameters(x, x_prime)
        beta_y, alpha_y, gamma_y, emittance_y = calculate_twiss_parameters(y, y_prime)
        
        # Calculate RMS beam size: σ_x = √(εβ)
        beam_size_x = np.sqrt(emittance_x * beta_x)
        beam_size_y = np.sqrt(emittance_y * beta_y)
        
        data_list.append({
            'freq': f_val,
            'beam_size_x': beam_size_x,
            'beam_size_y': beam_size_y,
            'emittance_x': emittance_x,
            'emittance_y': emittance_y,
            'beta_x': beta_x,
            'beta_y': beta_y,
            'alpha_x': alpha_x,
            'alpha_y': alpha_y
        })
    
    # Sort by frequency to match emmitance.py behavior
    data_list.sort(key=lambda x: x['freq'])
    
    # Extract sorted data
    for data in data_list:
        frequencies.append(data['freq'])
        x_beam_sizes.append(data['beam_size_x'])
        y_beam_sizes.append(data['beam_size_y'])
        x_emittances.append(data['emittance_x'])
        y_emittances.append(data['emittance_y'])
        x_beta_values.append(data['beta_x'])
        y_beta_values.append(data['beta_y'])
        x_alpha_values.append(data['alpha_x'])
        y_alpha_values.append(data['alpha_y'])
    
    # Create figure
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    axes = axes.flatten()
    
    # Convert to numpy arrays for easier handling
    frequencies = np.array(frequencies)
    x_beam_sizes = np.array(x_beam_sizes)
    y_beam_sizes = np.array(y_beam_sizes)
    x_emittances = np.array(x_emittances)
    y_emittances = np.array(y_emittances)
    x_beta_values = np.array(x_beta_values)
    y_beta_values = np.array(y_beta_values)
    x_alpha_values = np.array(x_alpha_values)
    y_alpha_values = np.array(y_alpha_values)
    
    # 1. Beam size vs frequency (LINEAR SCALE)
    scatter_x = axes[0].scatter(frequencies, x_beam_sizes, c=frequencies, cmap='plasma', s=50, alpha=0.7, label='X-direction')
    scatter_y = axes[0].scatter(frequencies, y_beam_sizes, c=frequencies, cmap='viridis', s=50, alpha=0.7, label='Y-direction')
    plt.colorbar(scatter_x, ax=axes[0], shrink=0.8, label='Frequency [Hz]')
    axes[0].set_xlabel('Frequency [Hz]')
    axes[0].set_ylabel('RMS Beam Size [mm] (linear scale)')
    axes[0].set_title('Beam Size vs Frequency')
    axes[0].grid(True, alpha=0.3)
    axes[0].legend()
    
    # 2. Emittance vs frequency (LINEAR SCALE)
    scatter_x = axes[1].scatter(frequencies, x_emittances, c=frequencies, cmap='plasma', s=50, alpha=0.7, label='X-direction')
    scatter_y = axes[1].scatter(frequencies, y_emittances, c=frequencies, cmap='viridis', s=50, alpha=0.7, label='Y-direction')
    plt.colorbar(scatter_x, ax=axes[1], shrink=0.8, label='Frequency [Hz]')
    axes[1].set_xlabel('Frequency [Hz]')
    axes[1].set_ylabel('Emittance [mm·rad] (linear scale)')
    axes[1].set_title('Emittance vs Frequency')
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    # 3. Beta function vs frequency (LINEAR SCALE)
    axes[2].scatter(frequencies, x_beta_values, c='blue', s=50, alpha=0.7, label='β_x')
    axes[2].scatter(frequencies, y_beta_values, c='red', s=50, alpha=0.7, label='β_y')
    axes[2].set_xlabel('Frequency [Hz]')
    axes[2].set_ylabel('Beta Function [mm/rad] (linear scale)')
    axes[2].set_title('Beta Function vs Frequency')
    axes[2].grid(True, alpha=0.3)
    axes[2].legend()
    
    # 4. Alpha parameter vs frequency (LINEAR SCALE)
    axes[3].scatter(frequencies, x_alpha_values, c='blue', s=50, alpha=0.7, label='α_x')
    axes[3].scatter(frequencies, y_alpha_values, c='red', s=50, alpha=0.7, label='α_y')
    axes[3].set_xlabel('Frequency [Hz]')
    axes[3].set_ylabel('Alpha Parameter')
    axes[3].set_title('Alpha Parameter vs Frequency')
    axes[3].grid(True, alpha=0.3)
    axes[3].axhline(y=0, color='k', linestyle='--', alpha=0.5)
    axes[3].legend()
    
    # 5. Beam size ratio (asymmetry) (LINEAR SCALE)
    size_ratio = np.array(x_beam_sizes) / np.array(y_beam_sizes)
    scatter_ratio = axes[4].scatter(frequencies, size_ratio, c=frequencies, cmap='plasma', s=50, alpha=0.7)
    plt.colorbar(scatter_ratio, ax=axes[4], shrink=0.8, label='Frequency [Hz]')
    axes[4].set_xlabel('Frequency [Hz]')
    axes[4].set_ylabel('X/Y Beam Size Ratio')
    axes[4].set_title('Beam Asymmetry (X/Y Ratio)')
    axes[4].grid(True, alpha=0.3)
    axes[4].axhline(y=1, color='r', linestyle='--', alpha=0.5, label='Symmetry line')
    axes[4].legend()
    
    # 6. Product of beam sizes (approximate envelope area) (LINEAR SCALE)
    envelope_area = np.array(x_beam_sizes) * np.array(y_beam_sizes)
    scatter_env = axes[5].scatter(frequencies, envelope_area, c=frequencies, cmap='plasma', s=50, alpha=0.7)
    plt.colorbar(scatter_env, ax=axes[5], shrink=0.8, label='Frequency [Hz]')
    axes[5].set_xlabel('Frequency [Hz]')
    axes[5].set_ylabel('X×Y Product [mm²] (linear scale)')
    axes[5].set_title('Beam Envelope Area (X×Y)')
    axes[5].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Beam size analysis plot (linear scale) saved to: {save_path}")
    plt.close(fig)
    
    # Log comprehensive statistics
    console.log_section("束流尺寸和 TWISS 参数分析 (线性尺度)")
    
    summary_rows = []
    for direction, sizes, emittances, betas, alphas, freqs in [
        ('X', x_beam_sizes, x_emittances, x_beta_values, x_alpha_values, frequencies),
        ('Y', y_beam_sizes, y_emittances, y_beta_values, y_alpha_values, frequencies)
    ]:
        min_idx = np.argmin(sizes)
        min_size = sizes[min_idx]
        freq_at_min = freqs[min_idx]
        emit_at_min = emittances[min_idx]
        beta_at_min = betas[min_idx]
        alpha_at_min = alphas[min_idx]
        
        # Determine convergence status
        if alpha_at_min > 0.01:
            status = "收敛"
        elif alpha_at_min < -0.01:
            status = "发散"
        else:
            status = "束腰"
        
        summary_rows.append([
            direction, f"{min_size:.2e}", str(freq_at_min), f"{emit_at_min:.2e}", 
            f"{beta_at_min:.2e}", f"{alpha_at_min:.2e}", status
        ])
    
    console.log_table(
        ['方向', '最小尺寸', '最小频率', '发射度', 'Beta', 'Alpha', '状态'],
        summary_rows
    )
    
    # Find frequency with minimum envelope area
    envelope_area = np.array(x_beam_sizes) * np.array(y_beam_sizes)
    min_area_idx = np.argmin(envelope_area)
    console.log(f"最小束流包络面积: {envelope_area[min_area_idx]:.2e} mm² 在频率 f={frequencies[min_area_idx]} Hz")
    
    # Find alpha local minima (points where alpha is close to zero)
    console.log_detailed_results("\n" + "=" * 60)
    console.log_detailed_results("Alpha 接近 0 的数据点 (局部最小值) - 线性尺度")
    console.log_detailed_results("=" * 60)
    
    def find_alpha_minima(freqs, alphas, direction_name):
        """Find local minima of |alpha| (points where alpha is close to zero)"""
        freqs = np.array(freqs)
        alphas = np.array(alphas)
        abs_alphas = np.abs(alphas)
        
        # Find local minima of |alpha| using find_peaks on -abs_alphas
        # prominence should be relative to the range of abs_alphas
        prominence = 0.01 * np.ptp(abs_alphas)  # 1% of the range
        distance = len(alphas) // 100  # Minimum distance between peaks (1% of data points)
        
        # Find peaks in -abs_alphas (which are minima in abs_alphas)
        minima_indices, properties = find_peaks(-abs_alphas, prominence=prominence, distance=distance)
        
        # Also include points where |alpha| < threshold
        alpha_threshold = 0.001
        near_zero_indices = np.where(abs_alphas < alpha_threshold)[0]
        
        # Combine and remove duplicates
        all_indices = set(minima_indices) | set(near_zero_indices)
        
        # Create list of (freq, alpha) pairs, but filter by threshold
        all_points = []
        for idx in all_indices:
            if abs_alphas[idx] < alpha_threshold:
                all_points.append((freqs[idx], alphas[idx]))
        
        # Sort by frequency
        all_points.sort(key=lambda x: x[0])
        
        return all_points
    
    # X direction
    x_zero_points = find_alpha_minima(frequencies, x_alpha_values, "X")
    if x_zero_points:
        console.log_detailed_results(f"\nX 方向 (共 {len(x_zero_points)} 个数据点):")
        console.log_detailed_results(f"{'频率 (Hz)':<20} {'Alpha':<15}")
        console.log_detailed_results("-" * 60)
        for freq, alpha in x_zero_points:
            console.log_detailed_results(f"{freq:<20.0f} {alpha:<15.6e}")
    else:
        console.log_detailed_results("\nX 方向: 未找到 Alpha 接近 0 的数据点")
    
    # Y direction
    y_zero_points = find_alpha_minima(frequencies, y_alpha_values, "Y")
    if y_zero_points:
        console.log_detailed_results(f"\nY 方向 (共 {len(y_zero_points)} 个数据点):")
        console.log_detailed_results(f"{'频率 (Hz)':<20} {'Alpha':<15}")
        console.log_detailed_results("-" * 60)
        for freq, alpha in y_zero_points:
            console.log_detailed_results(f"{freq:<20.0f} {alpha:<15.6e}")
    else:
        console.log_detailed_results("\nY 方向: 未找到 Alpha 接近 0 的数据点")
    
    console.log_detailed_results("=" * 60)


def comprehensive_beam_analysis(file_type='csv'):
    """
    Comprehensive beam analysis with file selection dialog
    
    Parameters:
    file_type: 'csv', 'excel', or 'json'
    """
    # Select and read files using dialog
    dfs, filenames = select_and_read_files(file_type)
    
    if len(dfs) == 0:
        print("未找到有效文件。退出分析。")
        return
    
    print(f"\n开始分析 {len(dfs)} 个文件...")
    
    # Run beam size analysis
    plot_beam_size_vs_frequency_log(dfs, filenames)

    print("\n分析完成！")

# Main execution
if __name__ == "__main__":
    # Clear console output file
    console.clear_file()
    
    # Run comprehensive analysis with file selection dialog
    comprehensive_beam_analysis(file_type='csv')